import cv2
import os
import numpy as np

class FaceTrainer:
    def __init__(self, data_dir='data'):
        self.data_dir = data_dir
        self.recognizer = cv2.face.LBPHFaceRecognizer_create()
        
    def get_images_and_labels(self):
        """从data目录加载所有用户的人脸图像和标签"""
        image_paths = []
        labels = []
        
        # 遍历data目录下的所有用户文件夹
        for user_dir in os.listdir(self.data_dir):
            user_path = os.path.join(self.data_dir, user_dir)
            if not os.path.isdir(user_path):
                continue
                
            # 直接使用文件夹名作为用户ID（格式为"姓名_职业"）
            user_id = user_dir
            user_path = os.path.join(self.data_dir, user_dir)
            
            # 遍历用户文件夹中的所有图片
            for img_name in os.listdir(user_path):
                if not img_name.endswith('.jpg'):
                    continue
                    
                img_path = os.path.join(user_path, img_name)
                image_paths.append((img_path, user_id))
        
        # 读取所有图片并创建标签
        faces = []
        labels = []
        label_dict = {}  # 用户名到数字标签的映射
        current_label = 0
        
        for img_path, user_id in image_paths:
            # 读取灰度图像
            img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
            if img is None:
                continue
                
            # 为每个用户分配一个数字标签
            if user_id not in label_dict:
                label_dict[user_id] = current_label
                current_label += 1
                
            faces.append(img)
            labels.append(label_dict[user_id])
            
        return faces, labels, label_dict

    def train_model(self, output_file='wenjian/trainer.yml'):
        """训练模型并保存"""
        print("开始训练人脸识别模型...")
        
        # 获取训练数据
        faces, labels, label_dict = self.get_images_and_labels()
        
        if not faces:
            print("错误：没有找到训练数据！")
            return False
            
        # 训练模型
        self.recognizer.train(faces, np.array(labels))
        
        # 保存模型
        self.recognizer.save(output_file)
        print(f"模型训练完成，已保存到 {output_file}")
        
        # 保存标签映射关系
        with open('label_mapping.txt', 'w', encoding='utf-8') as f:
            for name, label in label_dict.items():
                f.write(f"{label}:{name}\\n")
                
        return True

if __name__ == "__main__":
    trainer = FaceTrainer()
    if trainer.train_model():
        print("训练成功！")
    else:
        print("训练失败，请检查数据目录")